package cn.xiaolang.function.data_structure;

import java.util.Timer;
import java.util.TimerTask;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicLong;

/**
 * 限流算法-令牌桶
 * <p>
 * 令牌桶算法能够支持的瞬时流量最大为桶大小
 * <p>
 * <p>
 * 令牌桶算法的原理是系统以恒定的速率产生令牌，然后把令牌放到令牌桶中，令牌桶有一个容量，
 * 当令牌桶满了的时候，再向其中放令牌，多余令牌会被丢弃；当想要处理一个请求的时候，
 * 需要从令牌桶中取出一个令牌，如果此时令牌桶中没有令牌，那么则拒绝该请求。
 * <p>
 * 缺点：
 * 但是需要一个Timer线程，资源占用较重
 *
 * @author fuhw/Dean
 * @date 2020-04-10
 */
public final class FlowLimitByTokenBucket {

    // 桶的容量
    private int capacity;
    // 当前令牌数
    private AtomicInteger tokens = new AtomicInteger(0);
    // 每秒颁发的令牌数
    private int issueTokenRatePerSecond;


    private FlowLimitByTokenBucket(int capacity, int issueTokenRatePerSecond) {
        this.capacity = capacity;
        this.issueTokenRatePerSecond = issueTokenRatePerSecond;
        new Timer().scheduleAtFixedRate(new TimerTask() {
            @Override
            public void run() {
                if (tokens.getAndAdd(issueTokenRatePerSecond) > capacity) {
                    // 丢弃多余令牌
                    tokens.set(capacity);
                }
            }
        }, 1000, 1000);
    }

    public static FlowLimitByTokenBucket build(int capacity, int issueTokenRatePerSecond) {
        return new FlowLimitByTokenBucket(capacity, issueTokenRatePerSecond);
    }

    private boolean acquire() {
        return acquire(1);
    }

    private boolean acquire(int permits) {
        if (tokens.get() < permits) {
            return false;
        }
        tokens.set(tokens.get() - permits);
        return true;
    }

    public static void main(String[] args) {
        FlowLimitByTokenBucket limit = FlowLimitByTokenBucket.build(100, 100);
        AtomicLong successCount = new AtomicLong();
        AtomicLong totalCount = new AtomicLong();
        long start = System.currentTimeMillis();
        while (true) {
            if (limit.acquire()) {
                System.out.println("当前token数：" + limit.tokens.get());
                successCount.getAndIncrement();
            }
            totalCount.getAndIncrement();

            long consume = (System.currentTimeMillis() - start) / 1000L;
            long qps = consume == 0 ? totalCount.get() : totalCount.get() / consume;
            System.out.println("total:" + totalCount.get() +
                    ",success:" + successCount.get() +
                    ",consume:" + consume + "(s)" +
                    ",qps:" + qps + "/s");
        }
    }
}
